#!/usr/bin/env python
# -*- coding: utf-8 -*-

from PyQt5.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, QPushButton, 
                            QLabel, QComboBox, QGroupBox, QFormLayout, 
                            QSpinBox, QDoubleSpinBox, QCheckBox, QProgressBar,
                            QMessageBox, QFileDialog, QTextEdit, QTabWidget,
                            QLineEdit, QRadioButton, QButtonGroup, QScrollArea,
                            QDialog, QInputDialog, QApplication)
from PyQt5.QtCore import Qt, QThread, pyqtSignal, QTimer, QSize
from PyQt5.QtGui import QIcon, QFont, QTextCursor
import os
import sys
import json
import yaml

# 训练线程，使用QThread避免UI冻结
class TrainingThread(QThread):
    # 定义信号
    update_progress = pyqtSignal(int)
    update_log = pyqtSignal(str)
    training_finished = pyqtSignal(bool, str)  # 成功/失败，消息
    
    def __init__(self, dataset_path, model_path, epochs, batch_size, img_size, device, save_dir):
        super().__init__()
        self.dataset_path = dataset_path
        self.model_path = model_path
        self.model_name = os.path.basename(model_path) if os.path.exists(model_path) else model_path
        self.epochs = epochs
        self.batch_size = batch_size
        self.img_size = img_size
        self.device = device
        self.save_dir = save_dir
        self.stop_flag = False
        
    def run(self):
        try:
            # 检查是否能导入ultralytics库
            try:
                from ultralytics import YOLO
                import torch
            except ImportError:
                self.update_log.emit("错误: 未安装ultralytics或PyTorch。请安装: pip install ultralytics")
                self.training_finished.emit(False, "未安装必要的库")
                return
                
            # 导入必要模块
            import time
            
            # 更新日志
            if os.path.exists(self.model_path):
                self.update_log.emit(f"正在使用本地模型: {self.model_path}")
            else:
                self.update_log.emit(f"正在使用/下载模型: {self.model_path}")
                
            self.update_log.emit(f"数据集路径: {self.dataset_path}")
            self.update_log.emit(f"训练参数: epochs={self.epochs}, batch_size={self.batch_size}, img_size={self.img_size}")
            self.update_log.emit(f"设备: {self.device}")
            self.update_log.emit("正在初始化模型...")
            
            # 检查数据集YAML是否存在
            if not os.path.exists(self.dataset_path):
                self.update_log.emit(f"错误：数据集文件不存在: {self.dataset_path}")
                self.training_finished.emit(False, f"数据集文件不存在: {self.dataset_path}")
                return
                
            # 验证YAML格式
            try:
                with open(self.dataset_path, 'r', encoding='utf-8') as f:
                    yaml_data = yaml.safe_load(f)
                    if yaml_data is None:
                        self.update_log.emit(f"错误：数据集YAML文件为空或格式错误")
                        self.training_finished.emit(False, "数据集YAML文件为空或格式错误")
                        return
                    
                # 检查YAML必要字段
                required_fields = ['train', 'val', 'nc', 'names']
                missing_fields = [field for field in required_fields if field not in yaml_data]
                
                if missing_fields:
                    error_msg = f"数据集YAML缺少必要字段: {', '.join(missing_fields)}"
                    self.update_log.emit(f"错误：{error_msg}")
                    self.update_log.emit("YAML文件应包含以下字段: train(训练集路径), val(验证集路径), nc(类别数), names(类别名称列表)")
                    self.update_log.emit("请修复数据集YAML文件或使用应用提供的'修复数据集'功能")
                    self.training_finished.emit(False, error_msg)
                    return
                
                # 检查names字段类型
                names = yaml_data.get('names')
                if not isinstance(names, list) and not isinstance(names, dict):
                    self.update_log.emit(f"错误：YAML中的'names'字段应为列表或字典，当前类型: {type(names)}")
                    self.training_finished.emit(False, "数据集YAML格式错误：'names'字段类型不正确")
                    return
                
                # 检查train和val路径是否设置
                train_path = yaml_data.get('train')
                val_path = yaml_data.get('val')
                
                if not train_path or not val_path:
                    self.update_log.emit(f"错误：train或val路径未设置")
                    self.training_finished.emit(False, "数据集YAML中train或val路径未设置")
                    return
                
                # 打印数据集信息以供参考
                nc = yaml_data.get('nc')
                class_names = yaml_data.get('names')
                self.update_log.emit(f"数据集类别数: {nc}")
                if isinstance(class_names, list):
                    self.update_log.emit(f"类别名称: {', '.join(class_names)}")
                else:
                    self.update_log.emit(f"类别索引和名称: {class_names}")
                
                # 如果路径是相对路径，基于YAML文件目录解析
                yaml_dir = os.path.dirname(os.path.abspath(self.dataset_path))
                
                # 检查路径是否存在但不阻止训练
                def check_path_exists(path, path_type="路径"):
                    if not path:
                        return
                        
                    if not os.path.isabs(path):
                        full_path = os.path.join(yaml_dir, path)
                    else:
                        full_path = path
                    
                    if not os.path.exists(full_path):
                        self.update_log.emit(f"警告：{path_type}可能不存在: {full_path}")
                    else:
                        self.update_log.emit(f"确认{path_type}存在: {full_path}")
                    return full_path
                
                # 检查训练和验证数据路径
                train_full_path = check_path_exists(train_path, "训练数据路径")
                val_full_path = check_path_exists(val_path, "验证数据路径")
                
            except Exception as e:
                self.update_log.emit(f"警告：解析数据集YAML时出错: {str(e)}")
                self.update_log.emit("将尝试继续训练，但可能会失败")
            
            # 尝试多种训练方法
            success = False
            
            # 方法1：直接使用模型训练
            self.update_log.emit("开始训练...")
            
            # 构建训练参数
            train_args = {
                'data': self.dataset_path,
                'epochs': self.epochs,
                'batch': self.batch_size,
                'imgsz': self.img_size,
                'device': self.device,
                'project': self.save_dir,
                'name': 'train',
                'exist_ok': True
            }
            
            # 打印训练命令
            cmd_str = f"yolo train"
            for k, v in train_args.items():
                if k == 'exist_ok':
                    cmd_str += f" {k}"
                else:
                    cmd_str += f" {k}={v}"
            self.update_log.emit(f"执行命令: {cmd_str}")
            
            # 创建日志回调函数
            class LogCallback:
                def __init__(self, emit_func):
                    self.emit_func = emit_func
                    self.buffer = ""
                    
                def write(self, text):
                    if text:
                        # 将文本添加到缓冲区
                        self.buffer += text
                        
                        # 如果有完整的行，发送到UI
                        if '\n' in self.buffer:
                            lines = self.buffer.split('\n')
                            # 处理除最后一行外的所有完整行
                            for line in lines[:-1]:
                                if line.strip():  # 避免空行
                                    self.emit_func(line.strip())
                            # 保留最后一个不完整的行
                            self.buffer = lines[-1]
                        
                        # 如果缓冲区太长，也发送
                        elif len(self.buffer) > 100:
                            self.emit_func(self.buffer.strip())
                            self.buffer = ""
                    
                    return len(text)
                    
                def flush(self):
                    # 刷新时发送剩余缓冲区内容
                    if self.buffer.strip():
                        self.emit_func(self.buffer.strip())
                        self.buffer = ""
                    pass
            
            # 重定向标准输出和标准错误到日志
            import sys
            original_stdout = sys.stdout
            original_stderr = sys.stderr
            log_callback = LogCallback(self.update_log.emit)
            sys.stdout = log_callback
            sys.stderr = log_callback
            
            try:
                # 初始化模型
                model = YOLO(self.model_path)
                
                # 执行训练（使用标准参数，不添加复杂的回调机制）
                self.update_log.emit("开始模型训练...")
                results = model.train(**train_args)
                
                # 训练完成
                self.update_progress.emit(100)
                self.update_log.emit("训练完成！")
                self.update_log.emit(f"模型已保存到: {os.path.join(self.save_dir, 'train')}")
                self.training_finished.emit(True, "训练成功完成")
                success = True
                
            except Exception as e:
                import traceback
                error_trace = traceback.format_exc()
                self.update_log.emit(f"训练过程中出错: {str(e)}")
                self.update_log.emit(error_trace)
                self.training_finished.emit(False, str(e))
            finally:
                # 确保缓冲区中所有内容都已输出
                log_callback.flush()
                # 恢复标准输出和标准错误
                sys.stdout = original_stdout
                sys.stderr = original_stderr
            
        except Exception as e:
            # 处理可能的异常
            import traceback
            error_msg = traceback.format_exc()
            self.update_log.emit(f"训练过程中出错: {str(e)}")
            self.update_log.emit(error_msg)
            self.training_finished.emit(False, str(e))
    
    def stop(self):
        """停止训练进程"""
        self.stop_flag = True
        self.update_log.emit("已发送停止训练信号，等待当前epoch完成...")

class TrainingPanel(QWidget):
    """模型训练面板"""
    
    def __init__(self):
        super().__init__()
        self.init_ui()
        self.training_thread = None
        
    def init_ui(self):
        """初始化UI"""
        main_layout = QVBoxLayout(self)
        main_layout.setContentsMargins(20, 20, 20, 20)
        main_layout.setSpacing(20)
        
        # 创建滚动区域，允许在小窗口中查看所有内容
        scroll_area = QScrollArea()
        scroll_area.setWidgetResizable(True)
        scroll_area.setFrameShape(QScrollArea.NoFrame)
        
        scroll_widget = QWidget()
        scroll_layout = QVBoxLayout(scroll_widget)
        scroll_layout.setContentsMargins(0, 0, 0, 0)
        scroll_layout.setSpacing(20)
        
        # 标题标签
        title_label = QLabel("模型训练")
        title_label.setProperty("title", "true")
        title_label.setAlignment(Qt.AlignCenter)
        title_font = QFont()
        title_font.setPointSize(16)
        title_font.setBold(True)
        title_label.setFont(title_font)
        scroll_layout.addWidget(title_label)
        
        # 数据集设置
        dataset_group = QGroupBox("数据集设置")
        dataset_layout = QFormLayout(dataset_group)
        dataset_layout.setFieldGrowthPolicy(QFormLayout.ExpandingFieldsGrow)
        
        # 数据集路径选择
        dataset_path_layout = QHBoxLayout()
        self.dataset_path_edit = QLineEdit()
        self.dataset_path_edit.setReadOnly(True)
        self.dataset_path_edit.setPlaceholderText("选择项目数据集或YAML文件")
        
        dataset_browse_btn = QPushButton("浏览...")
        dataset_browse_btn.setFixedWidth(100)
        dataset_browse_btn.clicked.connect(self.browse_dataset)
        
        dataset_path_layout.addWidget(self.dataset_path_edit)
        dataset_path_layout.addWidget(dataset_browse_btn)
        
        dataset_layout.addRow("数据集路径:", dataset_path_layout)
        
        # 使用当前项目数据集
        use_project_dataset_btn = QPushButton("使用当前项目数据集")
        use_project_dataset_btn.clicked.connect(self.use_current_project_dataset)
        dataset_layout.addRow("", use_project_dataset_btn)
        
        scroll_layout.addWidget(dataset_group)
        
        # 模型设置
        model_group = QGroupBox("模型设置")
        model_layout = QFormLayout(model_group)
        
        # 选择预训练模型
        self.model_combo = QComboBox()
        
        # 刷新模型列表按钮
        refresh_model_btn = QPushButton("刷新列表")
        refresh_model_btn.setFixedWidth(80)
        refresh_model_btn.clicked.connect(self.scan_available_models)
        
        # 下载模型按钮
        download_model_btn = QPushButton("下载模型")
        download_model_btn.setFixedWidth(80)
        download_model_btn.clicked.connect(self.download_official_model)
        
        model_selection_layout = QHBoxLayout()
        model_selection_layout.addWidget(self.model_combo, 1)
        model_selection_layout.addWidget(refresh_model_btn)
        model_selection_layout.addWidget(download_model_btn)
        
        model_layout.addRow("预训练模型:", model_selection_layout)
        
        # 选择保存目录
        save_dir_layout = QHBoxLayout()
        self.save_dir_edit = QLineEdit()
        self.save_dir_edit.setReadOnly(True)
        self.save_dir_edit.setPlaceholderText("选择模型保存目录")
        
        save_dir_browse_btn = QPushButton("浏览...")
        save_dir_browse_btn.setFixedWidth(100)
        save_dir_browse_btn.clicked.connect(self.browse_save_dir)
        
        save_dir_layout.addWidget(self.save_dir_edit)
        save_dir_layout.addWidget(save_dir_browse_btn)
        
        model_layout.addRow("保存目录:", save_dir_layout)
        
        scroll_layout.addWidget(model_group)
        
        # 训练参数设置
        training_group = QGroupBox("训练参数")
        training_layout = QFormLayout(training_group)
        
        # 训练轮数
        self.epochs_spin = QSpinBox()
        self.epochs_spin.setRange(1, 1000)
        self.epochs_spin.setValue(50)
        training_layout.addRow("训练轮数:", self.epochs_spin)
        
        # 批次大小
        self.batch_size_combo = QComboBox()
        self.batch_size_combo.addItems(["4", "8", "16", "32", "64"])
        self.batch_size_combo.setCurrentText("16")
        training_layout.addRow("批次大小:", self.batch_size_combo)
        
        # 图像尺寸
        self.img_size_combo = QComboBox()
        self.img_size_combo.addItems(["320", "416", "512", "640", "768"])
        self.img_size_combo.setCurrentText("640")
        training_layout.addRow("图像尺寸:", self.img_size_combo)
        
        # 训练设备选择
        self.device_combo = QComboBox()
        
        # 检测系统中的CUDA设备
        try:
            import torch
            if torch.cuda.is_available():
                for i in range(torch.cuda.device_count()):
                    device_name = f"cuda:{i} ({torch.cuda.get_device_name(i)})"
                    self.device_combo.addItem(device_name)
            else:
                self.device_combo.addItem("cpu")
        except:
            self.device_combo.addItem("cpu")
            self.device_combo.addItem("cuda:0")  # 默认添加，即使可能不可用
            
        training_layout.addRow("训练设备:", self.device_combo)
        
        scroll_layout.addWidget(training_group)
        
        # 训练控制
        control_group = QGroupBox("训练控制")
        control_layout = QVBoxLayout(control_group)
        
        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setRange(0, 100)
        self.progress_bar.setValue(0)
        control_layout.addWidget(self.progress_bar)
        
        # 按钮布局
        buttons_layout = QHBoxLayout()
        
        self.start_btn = QPushButton("开始训练")
        self.start_btn.clicked.connect(self.start_training)
        
        self.stop_btn = QPushButton("停止训练")
        self.stop_btn.clicked.connect(self.stop_training)
        self.stop_btn.setEnabled(False)
        
        buttons_layout.addWidget(self.start_btn)
        buttons_layout.addWidget(self.stop_btn)
        
        control_layout.addLayout(buttons_layout)
        
        # 训练日志
        control_layout.addWidget(QLabel("训练日志:"))
        self.log_text = QTextEdit()
        self.log_text.setReadOnly(True)
        self.log_text.setMinimumHeight(200)
        control_layout.addWidget(self.log_text)
        
        scroll_layout.addWidget(control_group)
        
        # 设置滚动区域
        scroll_area.setWidget(scroll_widget)
        main_layout.addWidget(scroll_area)
        
        # 现在扫描已下载的模型
        self.scan_available_models()
        
    def browse_dataset(self):
        """浏览选择数据集路径"""
        dataset_path, _ = QFileDialog.getOpenFileName(
            self, "选择数据集YAML文件", "", "YAML文件 (*.yaml)"
        )
        if dataset_path:
            self.dataset_path_edit.setText(dataset_path)
            self.log_message(f"已选择数据集: {dataset_path}")
    
    def browse_save_dir(self):
        """浏览选择模型保存目录"""
        save_dir = QFileDialog.getExistingDirectory(
            self, "选择模型保存目录", ""
        )
        if save_dir:
            self.save_dir_edit.setText(save_dir)
            self.log_message(f"已选择保存目录: {save_dir}")
    
    def create_compatible_dataset_yaml(self, dataset_path):
        """创建兼容的数据集YAML文件
        
        如果数据集文件不兼容，尝试创建一个新的兼容版本
        返回新文件路径或原路径（如果没有问题）
        """
        try:
            with open(dataset_path, 'r', encoding='utf-8') as f:
                yaml_data = yaml.safe_load(f)
            
            # 检查YAML必要字段
            required_fields = ['train', 'val', 'nc', 'names']
            missing_fields = [field for field in required_fields if field not in yaml_data]
            
            if not missing_fields:
                return dataset_path  # 文件没有问题
                
            # 需要修复YAML文件
            self.log_message(f"尝试修复数据集YAML文件...")
            
            # 获取文件目录路径
            yaml_dir = os.path.dirname(os.path.abspath(dataset_path))
            
            # 创建基本结构
            fixed_yaml = {}
            
            # 复制现有字段
            for key, value in yaml_data.items():
                fixed_yaml[key] = value
                
            # 修复缺失字段
            if 'nc' not in fixed_yaml and 'names' in fixed_yaml:
                fixed_yaml['nc'] = len(fixed_yaml['names'])
                self.log_message(f"根据类别名称自动设置nc={fixed_yaml['nc']}")
                
            if 'names' not in fixed_yaml and 'nc' in fixed_yaml:
                # 创建默认类别名称
                fixed_yaml['names'] = [f'class{i}' for i in range(fixed_yaml['nc'])]
                self.log_message(f"自动创建默认类别名称")
                
            # 检查train和val路径
            if 'train' not in fixed_yaml or 'val' not in fixed_yaml:
                # 尝试推断数据路径
                images_dir = None
                if os.path.exists(os.path.join(yaml_dir, 'images')):
                    images_dir = os.path.join(yaml_dir, 'images')
                elif os.path.exists(os.path.join(yaml_dir, '..', 'images')):
                    images_dir = os.path.join(yaml_dir, '..', 'images')
                    
                train_dir = None
                val_dir = None
                
                if images_dir:
                    # 查找train和val子目录
                    if os.path.exists(os.path.join(images_dir, 'train')):
                        train_dir = os.path.join(images_dir, 'train')
                    else:
                        train_dir = images_dir
                        
                    if os.path.exists(os.path.join(images_dir, 'val')):
                        val_dir = os.path.join(images_dir, 'val')
                    else:
                        val_dir = images_dir
                
                # 设置相对路径
                if train_dir:
                    fixed_yaml['train'] = os.path.relpath(train_dir, yaml_dir)
                    self.log_message(f"自动设置train路径: {fixed_yaml['train']}")
                    
                if val_dir:
                    fixed_yaml['val'] = os.path.relpath(val_dir, yaml_dir)
                    self.log_message(f"自动设置val路径: {fixed_yaml['val']}")
            
            # 检查是否所有必要字段都已添加
            still_missing = [field for field in required_fields if field not in fixed_yaml]
            if still_missing:
                self.log_message(f"无法修复所有缺失字段: {still_missing}")
                return dataset_path
                
            # 创建修复后的YAML文件
            base_name = os.path.basename(dataset_path)
            fixed_name = f"fixed_{base_name}"
            fixed_path = os.path.join(yaml_dir, fixed_name)
            
            with open(fixed_path, 'w', encoding='utf-8') as f:
                yaml.dump(fixed_yaml, f, default_flow_style=False)
                
            self.log_message(f"已创建修复版数据集文件: {fixed_path}")
            return fixed_path
            
        except Exception as e:
            self.log_message(f"尝试修复数据集文件时出错: {str(e)}")
            return dataset_path
    
    def use_current_project_dataset(self):
        """使用当前项目的数据集"""
        try:
            # 获取主窗口
            main_window = self.window()
            if hasattr(main_window, 'current_project') and main_window.current_project:
                project_path = main_window.current_project.get('path')
                if project_path:
                    dataset_dir = os.path.join(project_path, "dataset")
                    dataset_yaml = os.path.join(dataset_dir, "dataset.yaml")
                    
                    if os.path.exists(dataset_yaml):
                        self.dataset_path_edit.setText(dataset_yaml)
                        self.log_message(f"已加载当前项目数据集: {dataset_yaml}")
                        
                        # 设置默认保存目录
                        models_dir = os.path.join(project_path, "models")
                        if not os.path.exists(models_dir):
                            os.makedirs(models_dir, exist_ok=True)
                        self.save_dir_edit.setText(models_dir)
                        
                        return
                        
                    # 如果没有找到yaml，则检查是否需要生成数据集
                    if os.path.exists(os.path.join(project_path, "images")) and os.path.exists(os.path.join(project_path, "labels")):
                        reply = QMessageBox.question(
                            self, 
                            "导出数据集", 
                            "项目中没有找到数据集配置，是否立即导出数据集？",
                            QMessageBox.Yes | QMessageBox.No
                        )
                        
                        if reply == QMessageBox.Yes:
                            # 通知主窗口导出数据集
                            if hasattr(main_window, 'export_current_project_dataset'):
                                main_window.export_current_project_dataset()
                                
                                # 导出后重新检查
                                if os.path.exists(dataset_yaml):
                                    self.dataset_path_edit.setText(dataset_yaml)
                                    self.log_message(f"已加载当前项目数据集: {dataset_yaml}")
                                    
                                    # 设置默认保存目录
                                    models_dir = os.path.join(project_path, "models")
                                    if not os.path.exists(models_dir):
                                        os.makedirs(models_dir, exist_ok=True)
                                    self.save_dir_edit.setText(models_dir)
                                    return
                        
                    self.log_message("未找到项目数据集，请先导出项目数据集")
                else:
                    self.log_message("无法获取当前项目路径")
            else:
                self.log_message("未打开项目，请先打开或创建一个项目")
        except Exception as e:
            self.log_message(f"获取项目数据集时出错: {str(e)}")
    
    def on_model_changed(self, index):
        """模型选择变更时的处理"""
        if index < 0:
            return
            
        model_name = self.model_combo.currentText()
        model_path = self.model_combo.itemData(index, Qt.ToolTipRole)
        
        if model_path:
            self.log_message(f"已选择模型: {model_name} ({model_path})")
        else:
            self.log_message(f"已选择模型: {model_name} (将从官方下载)")
    
    def start_training(self):
        """开始训练模型"""
        # 获取设置
        dataset_path = self.dataset_path_edit.text()
        model_index = self.model_combo.currentIndex()
        model_name = self.model_combo.currentText()
        model_path = self.model_combo.itemData(model_index, Qt.ToolTipRole) or model_name
        save_dir = self.save_dir_edit.text()
        epochs = self.epochs_spin.value()
        batch_size = int(self.batch_size_combo.currentText())
        img_size = int(self.img_size_combo.currentText())
        device = self.device_combo.currentText().split()[0]  # 只取设备ID部分
        
        # 验证输入
        if not dataset_path:
            QMessageBox.warning(self, "警告", "请选择数据集路径")
            return
            
        if not save_dir:
            QMessageBox.warning(self, "警告", "请选择模型保存目录")
            return
            
        # 确保必要的库已安装
        try:
            from ultralytics import YOLO
        except ImportError:
            reply = QMessageBox.question(
                self, 
                "安装依赖", 
                "未检测到Ultralytics库，需要安装才能进行训练。是否安装？",
                QMessageBox.Yes | QMessageBox.No
            )
            
            if reply == QMessageBox.Yes:
                self.log_message("请在命令行运行: pip install ultralytics")
                QMessageBox.information(
                    self,
                    "安装说明",
                    "请在命令行终端中执行以下命令：\n\n"
                    "pip install ultralytics\n\n"
                    "安装完成后重试。"
                )
            return
        
        # 如果模型没有完整路径且不是本地文件，检查是否需要下载
        if not os.path.exists(model_path):
            reply = QMessageBox.question(
                self, 
                "下载模型", 
                f"模型 {model_name} 不存在于本地，需要下载。是否继续？",
                QMessageBox.Yes | QMessageBox.No
            )
            
            if reply == QMessageBox.Yes:
                self.log_message(f"模型将在训练开始时自动下载: {model_name}")
            else:
                return
        
        # 检查并尝试修复数据集YAML文件
        if dataset_path.endswith('.yaml'):
            try:
                fixed_dataset_path = self.create_compatible_dataset_yaml(dataset_path)
                if fixed_dataset_path != dataset_path:
                    reply = QMessageBox.question(
                        self, 
                        "使用修复的数据集", 
                        f"检测到原始数据集格式可能有问题，已创建修复版本。\n\n是否使用修复后的数据集文件?\n\n{fixed_dataset_path}",
                        QMessageBox.Yes | QMessageBox.No
                    )
                    
                    if reply == QMessageBox.Yes:
                        dataset_path = fixed_dataset_path
                        self.log_message(f"使用修复后的数据集文件: {dataset_path}")
            except Exception as e:
                self.log_message(f"尝试检查/修复数据集文件时出错: {str(e)}")
        
        # 创建并启动训练线程
        self.training_thread = TrainingThread(
            dataset_path, model_path, epochs, batch_size, img_size, device, save_dir
        )
        self.training_thread.update_progress.connect(self.update_progress)
        self.training_thread.update_log.connect(self.log_message)
        self.training_thread.training_finished.connect(self.training_finished)
        
        # 更新UI状态
        self.start_btn.setEnabled(False)
        self.stop_btn.setEnabled(True)
        self.progress_bar.setValue(0)
        
        # 开始训练
        self.log_message(f"开始训练... 使用模型: {model_path}")
        self.training_thread.start()
    
    def stop_training(self):
        """停止训练"""
        if self.training_thread and self.training_thread.isRunning():
            reply = QMessageBox.question(
                self, 
                "停止训练", 
                "确定要停止训练吗？训练进程将被终止，可能导致当前epoch的数据丢失。",
                QMessageBox.Yes | QMessageBox.No
            )
            
            if reply == QMessageBox.Yes:
                # 更新UI以反映正在停止的状态
                self.stop_btn.setText("正在停止...")
                self.stop_btn.setEnabled(False)
                
                # 尝试通过标志先友好地通知停止
                self.training_thread.stop()
                self.log_message("已发送停止训练信号...")
                
                # 给线程一点时间来响应停止标志
                QTimer.singleShot(1000, self.force_stop_training)
    
    def force_stop_training(self):
        """强制停止训练线程"""
        if self.training_thread and self.training_thread.isRunning():
            # 如果线程仍在运行，使用terminate方法
            self.training_thread.terminate()
            self.log_message("已强制终止训练过程")
            
            # 恢复UI状态
            self.stop_btn.setText("停止训练")
            self.stop_btn.setEnabled(False)
            self.start_btn.setEnabled(True)
            
            # 通知用户
            self.log_message("训练已中止。如果要保存模型，可能需要在下一次训练时使用较低的学习率从头开始。")
    
    def update_progress(self, value):
        """更新进度条"""
        self.progress_bar.setValue(value)
    
    def log_message(self, message):
        """添加日志消息"""
        if not message:
            return
            
        # 根据消息类型设置不同颜色
        format_html = "<span style='color:{}'>{}</span>"
        
        if "错误" in message or "Error" in message or "exception" in message.lower():
            # 错误消息显示为红色
            formatted_message = format_html.format("#e74c3c", message)
        elif "完成" in message or "成功" in message or "已保存" in message:
            # 成功消息显示为绿色
            formatted_message = format_html.format("#2ecc71", message)
        elif "警告" in message or "Warning" in message:
            # 警告消息显示为黄色
            formatted_message = format_html.format("#f39c12", message)
        elif "epoch" in message.lower() or "val" in message.lower():
            # 训练进度消息显示为蓝色
            formatted_message = format_html.format("#3498db", message)
        else:
            # 普通消息
            formatted_message = message
            
        # 添加消息到日志窗口
        self.log_text.append(formatted_message)
        
        # 强制滚动到底部
        self.log_text.verticalScrollBar().setValue(self.log_text.verticalScrollBar().maximum())
        
        # 处理UI事件，确保日志及时显示
        QApplication.processEvents()
    
    def training_finished(self, success, message):
        """训练结束回调"""
        # 更新UI状态
        self.start_btn.setEnabled(True)
        self.stop_btn.setEnabled(False)
        
        if success:
            # 记录训练好的模型路径
            self.trained_model_path = os.path.join(self.save_dir_edit.text(), "train", "weights", "best.pt")
            
            if os.path.exists(self.trained_model_path):
                # 提示用户训练完成，并询问是否应用模型
                reply = QMessageBox.question(
                    self, 
                    "训练完成", 
                    f"模型训练成功！是否将训练好的模型复制到项目目录？\n\n模型路径: {self.trained_model_path}",
                    QMessageBox.Yes | QMessageBox.No
                )
                
                if reply == QMessageBox.Yes:
                    try:
                        # 获取主窗口
                        main_window = self.window()
                        if hasattr(main_window, 'current_project') and main_window.current_project:
                            project_path = main_window.current_project.get('path')
                            if project_path:
                                # 获取当前时间作为文件名一部分
                                import datetime
                                now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
                                model_filename = f"trained_model_{now}.pt"
                                
                                # 复制模型到项目目录
                                models_dir = os.path.join(project_path, "models")
                                if not os.path.exists(models_dir):
                                    os.makedirs(models_dir, exist_ok=True)
                                    
                                import shutil
                                dest_path = os.path.join(models_dir, model_filename)
                                shutil.copy2(self.trained_model_path, dest_path)
                                
                                self.log_message(f"模型已复制到项目目录: {dest_path}")
                                QMessageBox.information(self, "成功", f"模型已成功复制到项目目录。\n路径: {dest_path}")
                    except Exception as e:
                        self.log_message(f"复制模型到项目目录时出错: {str(e)}")
                        QMessageBox.warning(self, "错误", f"复制模型时出错: {str(e)}")
            else:
                QMessageBox.information(self, "训练完成", message)
        else:
            QMessageBox.warning(self, "训练失败", f"训练过程中出错: {message}")
            
        # 如果需要，可以在这里添加其他结束操作
        
    def scan_available_models(self):
        """扫描已下载的模型"""
        # 保存当前选择的模型（如果有）
        current_model = self.model_combo.currentText() if self.model_combo.count() > 0 else ""
        
        # 清空模型下拉框
        self.model_combo.clear()
        
        # 要扫描的目录列表
        scan_dirs = []
        
        # 1. 检查项目模型目录
        try:
            main_window = self.window()
            if hasattr(main_window, 'current_project') and main_window.current_project:
                project_path = main_window.current_project.get('path')
                if project_path:
                    project_models_dir = os.path.join(project_path, "models")
                    if os.path.exists(project_models_dir):
                        scan_dirs.append(project_models_dir)
                        self.log_message(f"扫描项目模型目录: {project_models_dir}")
        except Exception as e:
            self.log_message(f"扫描项目模型目录时出错: {str(e)}")
            
        # 2. 检查预训练模型目录
        pretrained_dir = "pretrained"
        if os.path.exists(pretrained_dir):
            scan_dirs.append(pretrained_dir)
            self.log_message(f"扫描预训练模型目录: {pretrained_dir}")
            
        # 3. 检查.ultralytics/models目录（ultralytics默认下载目录）
        home_dir = os.path.expanduser("~")
        ultralytics_models_dir = os.path.join(home_dir, ".ultralytics", "models")
        if os.path.exists(ultralytics_models_dir):
            scan_dirs.append(ultralytics_models_dir)
            self.log_message(f"扫描Ultralytics模型目录: {ultralytics_models_dir}")
            
        # 4. 检查models目录（可能的项目模型目录）
        if os.path.exists("models"):
            scan_dirs.append("models")
            self.log_message("扫描当前目录下的models目录")
            
        # 扫描所有目录并收集模型
        found_models = []
        for directory in scan_dirs:
            try:
                for file_name in os.listdir(directory):
                    file_path = os.path.join(directory, file_name)
                    if os.path.isfile(file_path) and file_name.endswith(".pt"):
                        # 记录完整路径和文件名
                        found_models.append((file_path, file_name))
            except Exception as e:
                self.log_message(f"扫描目录 {directory} 时出错: {str(e)}")
                
        # 添加找到的模型到下拉框
        for model_path, model_name in found_models:
            # 添加模型名称，在工具提示中显示完整路径
            self.model_combo.addItem(model_name)
            index = self.model_combo.count() - 1
            self.model_combo.setItemData(index, model_path, Qt.ToolTipRole)
        
        # 如果没有找到模型，添加默认YOLOv8模型
        if self.model_combo.count() == 0:
            default_models = ["yolov8n.pt", "yolov8s.pt", "yolov8m.pt", "yolov8l.pt", "yolov8x.pt"]
            for model_name in default_models:
                self.model_combo.addItem(model_name)
            self.log_message("未找到已下载的模型，将使用默认YOLOv8模型（需要联网下载）")
        else:
            self.log_message(f"找到 {self.model_combo.count()} 个模型文件")
            
        # 恢复之前选择的模型（如果可能）
        if current_model:
            index = self.model_combo.findText(current_model)
            if index >= 0:
                self.model_combo.setCurrentIndex(index)
        
        # 连接模型选择变更信号
        self.model_combo.currentIndexChanged.connect(self.on_model_changed)

    def download_official_model(self):
        """下载官方模型"""
        # 检查是否安装了ultralytics
        try:
            from ultralytics import YOLO
        except ImportError:
            QMessageBox.warning(self, "未安装依赖", "请先安装 ultralytics 库")
            return
        
        # 创建简单的选择对话框
        model_name, ok = QInputDialog.getItem(
            self, 
            "选择模型", 
            "选择要下载的模型:",
            ["yolov8n.pt", "yolov8s.pt", "yolov8m.pt", "yolov8l.pt", "yolov8x.pt"],
            0, 
            False
        )
        
        if not ok or not model_name:
            return
            
        # 选择保存目录
        # 尝试使用项目的models目录作为默认保存位置
        default_dir = ""
        try:
            main_window = self.window()
            if hasattr(main_window, 'current_project') and main_window.current_project:
                project_path = main_window.current_project.get('path')
                if project_path:
                    models_dir = os.path.join(project_path, "models")
                    if not os.path.exists(models_dir):
                        os.makedirs(models_dir, exist_ok=True)
                    default_dir = models_dir
        except:
            pass
            
        if not default_dir:
            if not os.path.exists("models"):
                os.makedirs("models", exist_ok=True)
            default_dir = os.path.abspath("models")
            
        save_dir = QFileDialog.getExistingDirectory(
            self, "选择保存目录", default_dir
        )
        
        if not save_dir:
            return
            
        # 开始下载
        self.log_message(f"开始下载模型: {model_name} 到 {save_dir}")
        
        try:
            # 创建模型实例（这会触发下载）
            model = YOLO(model_name)
            
            # 获取模型路径
            model_path = model.ckpt_path
            
            # 复制到指定目录
            import shutil
            dest_path = os.path.join(save_dir, os.path.basename(model_path))
            if model_path != dest_path:  # 如果不是同一个文件
                shutil.copy2(model_path, dest_path)
                
            self.log_message(f"模型已下载并保存到: {dest_path}")
            QMessageBox.information(self, "下载完成", f"模型已下载并保存到: {dest_path}")
            
            # 刷新模型列表
            self.scan_available_models()
        except Exception as e:
            self.log_message(f"下载模型失败: {str(e)}")
            QMessageBox.warning(self, "下载失败", f"下载模型失败: {str(e)}")

if __name__ == "__main__":
    import sys
    from PyQt5.QtWidgets import QApplication
    
    app = QApplication(sys.argv)
    window = TrainingPanel()
    window.show()
    sys.exit(app.exec_()) 